import torch
from torch import nn
import torch.nn.functional as F
from d2l import torch as d2l
import torchvision

# ResNet
# class Residual(nn.Module):
#     def __init__(self, input_channels,num_channels,use_1conv=False,strides=1):
#         super().__init__()
#         self.conv1 = nn.Conv2d(in_channels=input_channels,out_channels=num_channels,kernel_size=3,
#                                padding=1,stride=strides)
#         self.conv2 = nn.Conv2d(in_channels=num_channels,out_channels=num_channels,kernel_size=3,padding=1)
#         if use_1conv :
#             self.conv3 = nn.Conv2d(in_channels=input_channels,out_channels=num_channels,kernel_size=1,stride=strides)
#         else:
#             self.conv3 = None
#         self.bn1 = nn.BatchNorm2d(num_channels)
#         self.bn2 = nn.BatchNorm2d(num_channels)
#         self.relu = nn.ReLU(inplace=True)
#
#     def forward(self, x):
#         y = F.relu(self.bn1(self.conv1(x)))
#         y = self.bn2(self.conv2(y))
#         if self.conv3 is not None:
#             x = self.conv3(x)
#         y += x
#         return F.relu(y)
#
# b1 = nn.Sequential(nn.Conv2d(1,64,kernel_size=7,stride=2,padding=3,),
#                    nn.BatchNorm2d(64),nn.ReLU(),
#                    nn.MaxPool2d(kernel_size=3,stride=2,padding=1))
#
# def resnet_block(input_channels,num_channels,num_residuals,first_block=False):
#     blk = []
#     for i in range(num_residuals):
#         if i==0 and not first_block:
#             blk.append(Residual(input_channels,num_channels,use_1conv=True,strides=2))
#         else:
#             blk.append(Residual(num_channels,num_channels))
#     return blk
#
# b2 = nn.Sequential(*resnet_block(64,64,2,first_block=True))
# b3 = nn.Sequential(*resnet_block(64,128,2))
# b4 = nn.Sequential(*resnet_block(128,256,2))
# b5 = nn.Sequential(*resnet_block(256,512,2))
#
# net = nn.Sequential(b1,b2,b3,b4,b5,nn.AdaptiveAvgPool2d((1,1)),nn.Flatten(),nn.Linear(512,10))
#
# # x = torch.rand(1,1,224,224)
# # for layer in net:
# #     x = layer(x)
# #     print(f"{layer.__class__.__name__}:output shape{x.shape}")
#
# batch_size = 256
# lr = 0.05
# num_epochs = 10
# train_iter ,test_iter = d2l.load_data_fashion_mnist(batch_size,resize=224)
# d2l.train_ch6(net,train_iter,test_iter,num_epochs,lr,d2l.try_gpu())
# d2l.plt.show()

# d2l.set_figsize()
# img = d2l.Image.open('D:\\PytorchLearn\\pytorch_learning\\Project1\\capture_20221115110104367.bmp')
# d2l.plt.imshow(img)
# d2l.plt.show()
#
# def apply(img,aug,num_rows=2,num_cols=4,scale=1.5):
#     Y = [aug(img) for _ in range(num_rows*num_cols)]
#     d2l.show_images(Y,num_rows,num_cols,scale=scale)
#     d2l.plt.show()
#
# # apply(img,torchvision.transforms.RandomHorizontalFlip())
# # apply(img,torchvision.transforms.RandomVerticalFlip())
#
# shape_aug = torchvision.transforms.RandomResizedCrop(
#     (200,200),scale=(0.1,1),ratio=(1,1))
# apply(img,shape_aug)
#
# apply(img,torchvision.transforms.ColorJitter(brightness=0.5,contrast=0,saturation=0,hue=0))
#
# apply(img,torchvision.transforms.ColorJitter(brightness=0.5
#             ,contrast=0.5,saturation=0.5,hue=0.5))

all_images = torchvision.datasets.CIFAR10(root='./dataset', train=True, download=True)
# d2l.show_images([all_images[i][0] for i in range(32)],4,8,scale=0.8)
# d2l.plt.show()

train_augs = torchvision.transforms.Compose(
    [torchvision.transforms.RandomHorizontalFlip(),
    torchvision.transforms.ToTensor()])

test_augs = torchvision.transforms.Compose(
    [torchvision.transforms.ToTensor()])

def load_cifar10(is_train,augs,batch_size):
    dataset = torchvision.datasets.CIFAR10(root='./dataset', train=is_train,transform=augs, download=True)
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size,shuffle = is_train,
                        num_workers=d2l.get_dataloader_workers())
    return dataloader

def train_batch_ch13(net,X,y,loss,trainer,devices):
    if isinstance(X,list):
        X =[x.to(devices[0]) for x in X]
    else:
        X = X.to(devices[0])
    y = y.to(devices[0])
    net.train()
    trainer.zero_grad()
    pred = net(X)
    l = loss(pred, y)
    l.sum().backward()
    trainer.step()
    train_loss_sum = l.sum()
    train_acc_sum = d2l.accuracy(pred,y)
    return train_loss_sum, train_acc_sum

def train_ch13(net, train_iter, test_iter, loss, trainer, num_epochs,
               devices=d2l.try_all_gpus()):
    """用多GPU进行模型训练"""
    timer, num_batches = d2l.Timer(), len(train_iter)
    animator = d2l.Animator(xlabel='epoch', xlim=[1, num_epochs], ylim=[0, 1],
                            legend=['train loss', 'train acc', 'test acc'])
    net = nn.DataParallel(net, device_ids=devices).to(devices[0])
    for epoch in range(num_epochs):
        # 4个维度：储存训练损失，训练准确度，实例数，特点数
        metric = d2l.Accumulator(4)
        for i, (features, labels) in enumerate(train_iter):
            timer.start()
            l, acc = train_batch_ch13(
                net, features, labels, loss, trainer, devices)
            metric.add(l, acc, labels.shape[0], labels.numel())
            timer.stop()
            if (i + 1) % (num_batches // 5) == 0 or i == num_batches - 1:
                animator.add(epoch + (i + 1) / num_batches,
                             (metric[0] / metric[2], metric[1] / metric[3],
                              None))
        test_acc = d2l.evaluate_accuracy_gpu(net, test_iter)
        animator.add(epoch + 1, (None, None, test_acc))
    print(f'loss {metric[0] / metric[2]:.3f}, train acc '
          f'{metric[1] / metric[3]:.3f}, test acc {test_acc:.3f}')
    print(f'{metric[2] * num_epochs / timer.sum():.1f} examples/sec on '
          f'{str(devices)}')

batch_size=256
devices = d2l.try_all_gpus()
net = d2l.resnet18(10,3)

def init_weights(m):
    if type(m) in [nn.Linear,nn.Conv2d]:
        nn.init.xavier_uniform_(m.weight)

net.apply(init_weights)

def train_with_data_aug(train_augs,test_augs,net,lr=0.001):
    train_iter = load_cifar10(is_train=True,augs=train_augs,batch_size=batch_size)
    test_iter = load_cifar10(is_train=False,augs=test_augs,batch_size=batch_size)
    loss = nn.CrossEntropyLoss(reduction="none")
    trainer = torch.optim.Adam(net.parameters(), lr=lr)
    train_ch13(net,train_iter,test_iter,loss,trainer,15,devices)

train_with_data_aug(train_augs,test_augs,net,lr=0.001)






